import cv2
import numpy as np
import os
from glob import glob


def validate_masks(mask_dir, num_classes, mask_ext='.png'):
    """验证掩码文件的值范围"""
    mask_files = glob(os.path.join(mask_dir, '*' + mask_ext))

    print(f"检查 {len(mask_files)} 个掩码文件...")

    invalid_masks = []
    for mask_file in mask_files:
        mask = cv2.imread(mask_file, cv2.IMREAD_GRAYSCALE)
        unique_vals = np.unique(mask)
        min_val, max_val = mask.min(), mask.max()

        if min_val < 0 or max_val >= num_classes:
            invalid_masks.append((os.path.basename(mask_file), min_val, max_val, unique_vals))
            print(f"无效: {os.path.basename(mask_file)} - 范围 [{min_val}, {max_val}], 唯一值: {unique_vals}")

    if invalid_masks:
        print(f"\n发现 {len(invalid_masks)} 个无效掩码文件")
        return invalid_masks
    else:
        print("所有掩码文件值范围正确!")
        return None


def fix_masks(mask_dir, num_classes, mask_ext='.png'):
    """修复掩码文件的值范围"""
    mask_files = glob(os.path.join(mask_dir, '*' + mask_ext))

    for mask_file in mask_files:
        mask = cv2.imread(mask_file, cv2.IMREAD_GRAYSCALE)

        # 确保值在有效范围内
        mask = np.clip(mask, 0, num_classes - 1)

        # 保存修复后的掩码
        cv2.imwrite(mask_file, mask)

    print(f"已修复 {len(mask_files)} 个掩码文件")


if __name__ == '__main__':
    mask_dir = 'inputs/loveda_rural/masks'
    num_classes = 7

    # 验证掩码
    invalid_masks = validate_masks(mask_dir, num_classes)

    # 如果需要修复
    if invalid_masks:
        print("\n开始修复掩码文件...")
        fix_masks(mask_dir, num_classes)
        print("修复完成! 重新验证...")
        validate_masks(mask_dir, num_classes)